import numpy as np
import math
import matplotlib.pyplot as plt

# 频域抽取的基3FFT
N = 27
fs = 0.5
n = [2 * np.pi * fs * t / N for t in range(N)]
x = [int(np.round(np.sin(i) * 1024)) + int(np.round(np.cos(i) * 1024)) * 1j for i in n]
loop_num = int(round(math.log(N, 3)))
data = np.zeros((loop_num + 1, N), dtype=np.complex128)
data[0, :] = x

for i in range(loop_num):
    k = i + 1
    for p in range(3 ** i):
        for j in range(N // (3 ** k)):
            data[i + 1][j + 3 * p * (N // (3 ** k))] = data[i, j + 3 * p * (N // (3 ** k))] \
                                                       + data[i, j + N // (3 ** k) + 3 * p * (N // (3 ** k))] \
                                                       + data[i, j + 2 * N // (3 ** k) + 3 * p * (N // (3 ** k))]
            data[i + 1][j + N // (3 ** k) + 3 * p * (N // (3 ** k))] = (data[i, j + 3 * p * (N // (3 ** k))] \
                                                                        + data[i, j + N // (3 ** k) + 3 * p * (
                            N // (3 ** k))] * np.e ** (-1j * 2 * np.pi / 3) \
                                                                        + data[i, j + 2 * N // (3 ** k) + 3 * p * (
                            N // (3 ** k))] * np.e ** (-1j * 4 * np.pi / 3)) * np.e ** (
                                                                               -1j * 2 * j * np.pi * (3 ** i) / N)
            data[i + 1][j + 2 * N // (3 ** k) + 3 * p * (N // (3 ** k))] = (data[i, j + 3 * p * (N // (3 ** k))]
                                                                            + data[i, j + N // (3 ** k) + 3 * p * (
                            N // (3 ** k))] * np.e ** (-1j * 4 * np.pi / 3)
                                                                            + data[i, j + 2 * N // (3 ** k) + 3 * p * (
                            N // (3 ** k))] * np.e ** (-1j * 2 * np.pi / 3)) * np.e ** (-1j * 4 * j * np.pi * (
                    3 ** i) / N)


# 递推计算倒序，
def rev3(k, N):
    if (k == 0):
        return (0)
    else:
        return (((rev3(k // 3, N) // 3) + (k % 3) * (N // 3)))


# 输出data的值
print("FFT结果：")
# print(data)
# 输出倒序
fft_out = np.ones_like(data[0, :])
for k in range(N):
    fft_out[rev3(k, N)] = data[loop_num, k]

print(fft_out)

xf = np.fft.fft(x, N)
plt.plot(abs(xf - fft_out))
plt.xlabel('Frequency')
plt.ylabel('Amplitude Error')
plt.title('Amplitude Error between numpy FFT and mixed radix FFT')
plt.show()

fft3_cmp = xf - fft_out
plt.plot(fft3_cmp.imag)
plt.plot(fft3_cmp.real)
